'''
Desc: YOLO+PyQt5界面设计
Author: 陈爱国
Date: 2024-09-17 20:35:20
'''
from PyQt5.QtWidgets import *
from PyQt5.QtCore import QTimer
from PyQt5.QtGui import QImage, QPixmap
import cv2, os, time
from ultralytics import YOLO
from threading import Thread

# 设置YOLO处理不要输出调试信息
os.environ['YOLO_VERBOSE'] = 'False'

class Yolo(QMainWindow):
    def __init__(self):
        super(Yolo, self).__init__()
        self.initUi()

        # 点击摄像头处理
        self.btnCamera.clicked.connect(self.startCameraEvent)
        self.btnStop.clicked.connect(self.stopBtnEvent)

        # 定义定时器
        self.time_camera = QTimer()
        self.time_camera.timeout.connect(self.show_camera)

        # 加载YOLO nano模型，第一次比较耗时，要20s左右
        self.model = YOLO('./pt/yolov8n.pt')

        # 要处理的视频帧队列，目前就存放1帧图片
        self.frameToAnalyze=[]

        # 启动处理视频帧独立线程
        Thread(target=self.frameAnalyzeThreadFunc, daemon=True).start()

    def initUi(self):
        # 设置窗口标题
        self.setWindowTitle("YOLO+PyQt5实例")
        # 设置窗口的最小尺寸
        self.setMinimumSize(1000, 700)
        # 窗口最大化显示
        self.showMaximized()

        # centralwidget
        centralWidget = QWidget(self)
        self.setCentralWidget(centralWidget)

        # centralwidget中的主布局
        mainLayout = QVBoxLayout(centralWidget)

        # 初始化top: 图形展示
        self.initTop()
        mainLayout.addLayout(self.topLayout)

        # 初始化bottom
        self.initBottom()
        mainLayout.addWidget(self.groupbox)

    def initTop(self):
        '''
        Desc: 初始化上半部分
        Author: 陈爱国
        Date: 2024-09-17 21:00:30
        '''        
        self.topLayout = QHBoxLayout()
        self.lbl_ori_video = QLabel()
        self.lbl_treated_video = QLabel()
        self.lbl_ori_video.setMinimumSize(480, 480)
        self.lbl_treated_video.setMinimumSize(480, 480)
        self.lbl_ori_video.setStyleSheet('border: 1px solid #D7E2F9;')
        self.lbl_treated_video.setStyleSheet('border: 1px solid #D7E2F9;')

        self.topLayout.addWidget(self.lbl_ori_video)
        self.topLayout.addWidget(self.lbl_treated_video)

    def initBottom(self):
        '''
        Desc: 初始化下半部分
        Author: 陈爱国
        Date: 2024-09-17 21:00:56
        '''        
        self.groupbox = QGroupBox()

        bottomLayout = QHBoxLayout()
        self.groupbox.setLayout(bottomLayout)

        self.textLog = QTextBrowser()

        btnLayout = QVBoxLayout()
        self.btnVideo = QPushButton("视频文件")
        self.btnCamera = QPushButton("摄像头")
        self.btnStop = QPushButton("停止")

        btnLayout.addWidget(self.btnVideo)
        btnLayout.addWidget(self.btnCamera)
        btnLayout.addWidget(self.btnStop)

        bottomLayout.addWidget(self.textLog)
        bottomLayout.addLayout(btnLayout)

    def startCameraEvent(self):
        '''
        Desc: 摄像头按钮事件处理
        Author: 陈爱国
        Date: 2024-09-18 21:23:41
        '''        
        # 初始化摄像头 cv2.CAP_DSHOW可以让摄像头的打开速度比较块
        self.capture = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        if not self.capture.isOpened():
            print("摄像头不能打开")
            return

        # 启动定时器
        if not self.time_camera.isActive():
            self.time_camera.start(50)

    def show_camera(self):
        '''
        Desc: 显示摄像头图片
        Author: 陈爱国
        Date: 2024-09-18 21:43:52
        '''        
        # 从视频中读取一帧
        ret, frame = self.capture.read()
        # 未读取成功
        if not ret:
            return
        
        # 把读取到的16:10的帧大小重新设置
        frame = cv2.resize(frame, (520, 400))
        # 视频色彩转回RGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 变成QImage形式
        qImage = QImage(frame.data, frame.shape[1], frame.shape[0],
                        QImage.Format_RGB888)
        # 将图片设置到label上
        self.lbl_ori_video.setPixmap(QPixmap.fromImage(qImage))

        # 如果当前没有处理任务
        if not self.frameToAnalyze:
            self.frameToAnalyze.append(frame)

    def frameAnalyzeThreadFunc(self):
        while True:
            if not self.frameToAnalyze:
                time.sleep(0.01)
                continue

            frame = self.frameToAnalyze.pop(0)
            results = self.model(frame)[0]
            img = results.plot(line_width=1)
            # 变成QImage形式
            qImage = QImage(img.data, img.shape[1], img.shape[0],
                            QImage.Format_RGB888)
            self.lbl_treated_video.setPixmap(QPixmap.fromImage(qImage))

            time.sleep(0.5)


    def stopBtnEvent(self):
        # 停止定时器
        self.time_camera.stop()

        # 关闭摄像头
        self.capture.release()

        # 清空图片
        self.lbl_ori_video.clear()
        self.lbl_treated_video.clear()

        


